package org.xyl.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.xyl.service.VectorLibraryService;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.data.segment.TextSegment;

import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/api/vector-library")
public class VectorLibraryController {

    @Autowired
    private VectorLibraryService vectorLibraryService;

    /**
     * 创建新的向量库
     */
    @PostMapping("/create")
    public ResponseEntity<?> createLibrary(@RequestParam("libraryName") String libraryName) {
        boolean created = vectorLibraryService.createVectorLibrary(libraryName);
        if (created) {
            return ResponseEntity.ok(Map.of("message", "向量库创建成功", "libraryName", libraryName));
        } else {
            return ResponseEntity.badRequest().body(Map.of("message", "向量库已存在"));
        }
    }

    /**
     * 删除向量库
     */
    @DeleteMapping("/delete/{libraryName}")
    public ResponseEntity<?> deleteLibrary(@PathVariable("libraryName") String libraryName) {
        boolean deleted = vectorLibraryService.deleteVectorLibrary(libraryName);
        if (deleted) {
            return ResponseEntity.ok(Map.of("message", "向量库删除成功"));
        } else {
            return ResponseEntity.notFound().build();
        }
    }

    /**
     * 列出所有向量库
     */
    @GetMapping("/list")
    public ResponseEntity<List<String>> listLibraries() {
        List<String> libraries = vectorLibraryService.listVectorLibraries();
        return ResponseEntity.ok(libraries);
    }

    /**
     * 获取向量库元数据
     */
    @GetMapping("/metadata/{libraryName}")
    public ResponseEntity<?> getLibraryMetadata(@PathVariable("libraryName") String libraryName) {
        VectorLibraryService.VectorLibraryMetadata metadata = vectorLibraryService.getLibraryMetadata(libraryName);
        if (metadata != null) {
            return ResponseEntity.ok(metadata);
        } else {
            return ResponseEntity.notFound().build();
        }
    }

    /**
     * 向向量库添加文档
     */
    @PostMapping("/{libraryName}/documents")
    public ResponseEntity<?> addDocument(@PathVariable("libraryName") String libraryName,
                                        @RequestBody Map<String, Object> requestBody) {
        String documentText = (String) requestBody.get("text");
        
        // 获取元数据，确保是Object类型
        @SuppressWarnings("unchecked")
        Map<String, Object> metadata = (Map<String, Object>) requestBody.getOrDefault("metadata", Map.of());

        if (documentText == null || documentText.isEmpty()) {
            return ResponseEntity.badRequest().body(Map.of("message", "文档文本不能为空"));
        }

        // 调用服务方法添加文档
        try {
            vectorLibraryService.addDocumentToLibrary(libraryName, documentText, metadata);
            return ResponseEntity.ok(Map.of("message", "文档添加成功"));
        } catch (Exception e) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(Map.of("message", "文档添加失败: " + e.getMessage()));
        }
    }

    /**
     * 从向量库搜索相关文档
     */
    @GetMapping("/{libraryName}/search")
    public ResponseEntity<?> searchInLibrary(@PathVariable("libraryName") String libraryName,
                                           @RequestParam("query") String query,
                                           @RequestParam(name ="maxResults", defaultValue = "5") int maxResults) {
        List<EmbeddingMatch<TextSegment>> results = vectorLibraryService.searchInLibrary(libraryName, query, maxResults);
        
        // 转换结果为更友好的格式
        List<Map<String, Object>> formattedResults = new java.util.ArrayList<>();
        for (EmbeddingMatch<TextSegment> match : results) {
            Map<String, Object> result = new java.util.HashMap<>();
            result.put("text", match.embedded().text());
            result.put("score", match.score());
            formattedResults.add(result);
        }
        
        return ResponseEntity.ok(Map.of(
                "query", query,
                "results", formattedResults,
                "total", results.size()
        ));
    }
}